#========================================
# author: changlong.zang
#   mail: zclongpop@163.com
#   date: Wed, 06 Jul 2016, 17:03:11
#========================================
import re, os, time, mmap, itertools, getpass, pymel.core
import maya.cmds             as mc
import maya.OpenMaya         as OpenMaya
import maya.OpenMayaUI       as OpenMayaUI
import xml.etree.ElementTree as etTree
import xml.dom.minidom       as minidom
from foleyUtils              import muuid, scriptTool, mayaTool
#--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#=========================================
# Utils
#=========================================

def get_input_dg_nodes(mobject):
    '''
    '''
    iterator = OpenMaya.MItDependencyGraph(mobject, OpenMaya.MFn.kDependencyNode, OpenMaya.MItDependencyGraph.kUpstream)
    while not iterator.isDone():
        yield iterator.currentItem()
        iterator.next()





#=========================================
# Get and Export node and data
#=========================================

def get_scene_SG_nodes():
    '''
    '''
    iterator = OpenMaya.MItDependencyNodes(OpenMaya.MFn.kShadingEngine)
    while not iterator.isDone():
        if not OpenMaya.MFnDependencyNode(iterator.item()).isDefaultNode():
            yield iterator.item()
        iterator.next()





def get_SG_connect_geometry(mobject):
    '''
    '''
    MFn_Set   = OpenMaya.MFnSet(mobject)
    selection = OpenMaya.MSelectionList()
    MFn_Set.getMembers(selection, False)

    return selection





def get_geometry_connect_SG(mobject):
    '''
    '''
    iterator = OpenMaya.MItDependencyGraph(mobject,
                                           OpenMaya.MFn.kShadingEngine,
                                           OpenMaya.MItDependencyGraph.kDownstream)

    while not iterator.isDone():
        if not OpenMaya.MFnDependencyNode(iterator.currentItem()).isDefaultNode():
            yield iterator.currentItem()
        iterator.next()





def get_select_SG_nodes():
    '''
    '''
    selection = OpenMaya.MSelectionList()
    OpenMaya.MGlobal.getActiveSelectionList(selection)

    sg_list  = OpenMaya.MSelectionList()

    iterator = OpenMaya.MItSelectionList(selection)
    mobject  = OpenMaya.MObject()
    while not iterator.isDone():
        iterator.getDependNode(mobject)

        for sg in get_geometry_connect_SG(mobject):
            if not sg_list.hasItem(sg):
                sg_list.add(sg)
                yield sg

        iterator.next()





def get_SG_xml_node(mobject):
    '''
    '''
    return etTree.Element('shadingEngine', name=muuid.get_scene_object_name(mobject), uuid=muuid.get_uuid_by_object(mobject))





def get_geometry_xml_nodes(selection):
    '''
    '''
    iterator = OpenMaya.MItSelectionList(selection)

    mobject  = OpenMaya.MObject()
    objects  = list()
    while not iterator.isDone():
        iterator.getStrings(objects)

        iterator.getDependNode(mobject)
        uuid_hex = muuid.get_uuid_by_object(mobject)

        for obj in objects:
            yield etTree.Element('geometry', name=re.split('[\|:](?!\d+\])', obj)[-1], uuid=uuid_hex)

        iterator.next()





def get_SG_connect_data(mobject):
    '''
    '''
    sg_connected_geometry = get_SG_connect_geometry(mobject)
    if sg_connected_geometry.isEmpty():
        return None

    SG_xml_node = get_SG_xml_node(mobject)

    for geo_xml_node in get_geometry_xml_nodes(sg_connected_geometry):
        SG_xml_node.append(geo_xml_node)

    return SG_xml_node





def get_SG_nodes_connect_data(selection):
    '''
    '''
    scene_xml_data  = etTree.Element('Shading-Data', artist=getpass.getuser(), time=time.strftime("%a, %d %b %Y, %H:%M:%S", time.localtime()))

    for mobject in selection:
        sg_xml_node = get_SG_connect_data(mobject)
        if sg_xml_node is not None:
            scene_xml_data.append(sg_xml_node)

    return scene_xml_data





def export_SG_connect_data(xml_data, filePath):
    '''
    '''
    if not filePath:
        return

    xml_tree = etTree.ElementTree(xml_data)
    try:
        xml_tree.write(filePath)
    except:
        return False





def export_all_SG_connect_data(filePath):
    '''
    '''
    scene_xml_data = get_SG_nodes_connect_data(get_scene_SG_nodes())
    export_SG_connect_data(scene_xml_data, filePath)

    return True





def export_sel_SG_connect_data(filePath):
    '''
    '''
    scene_xml_data = get_SG_nodes_connect_data(get_select_SG_nodes())
    export_SG_connect_data(scene_xml_data, filePath)

    return True





def export_SG_nodes(nodes, filePath):
    '''
    '''
    if not filePath:
        return

    selection = OpenMaya.MSelectionList()
    for sg in nodes:
        if get_SG_connect_geometry(sg).isEmpty():
            continue
        selection.add(sg)

    if selection.isEmpty():
        return

    OpenMaya.MGlobal.setActiveSelectionList(selection)
    OpenMaya.MFileIO.exportSelected(filePath, None, True)





def export_all_SG_nodes(filePath):
    '''
    '''
    export_SG_nodes(get_scene_SG_nodes(), filePath)





def export_sel_SG_nodes(filePath):
    '''
    '''
    export_SG_nodes(get_select_SG_nodes(), filePath)





#=========================================
# Set and Import node and data
#=========================================
def import_objects_keep_select_wraper(func):
    '''
    '''
    def doIt(*args, **kargs):
        '''
        '''
        selection = OpenMaya.MSelectionList()
        OpenMaya.MGlobal.getActiveSelectionList(selection)

        result = func(*args, **kargs)

        OpenMaya.MGlobal.setActiveSelectionList(selection)

        return result

    doIt.__name__ = func.__name__
    doIt.__doc__  = func.__doc__

    return doIt






@import_objects_keep_select_wraper
def import_SG_nodes(filePath):
    '''
    '''
    if not os.path.isfile(filePath):
        return

    filePath = filePath.replace('\\', '/')

    #- reference shader file...
    OpenMaya.MFileIO.reference(filePath, False, False, 'shader')

    #- get scene shadingEngine nodes with reference...
    scene_sg_nodes_uuid = [x for x, y in muuid.get_scene_shadingEngine_uuid_data(False)]

    #- list reference file nodes
    reference_nodes = OpenMaya.MSelectionList()
    OpenMaya.MFileIO.getReferenceNodes(filePath, reference_nodes)

    #- init reference all shadingEngine nodes and scene haven't shading engine nodes...
    touse_sg_nodes = OpenMaya.MSelectionList()
    unuse_sg_nodes = OpenMaya.MSelectionList()

    iterator = OpenMaya.MItSelectionList(reference_nodes)
    mobject  = OpenMaya.MObject()
    while not iterator.isDone():
        iterator.getDependNode(mobject)
        #- just save SG nodes..
        if mobject.apiType() != OpenMaya.MFn.kShadingEngine:
            iterator.next()
            continue

        #- 
        if muuid.get_uuid_by_object(mobject) in scene_sg_nodes_uuid:
            unuse_sg_nodes.add(mobject)
        else:
            touse_sg_nodes.add(mobject)

        iterator.next()

    #- remove reference shanding file...
    if touse_sg_nodes.isEmpty():
        OpenMaya.MFileIO.removeReference(filePath)

    #- import referenced shader file, and delete duplicates shadingEngine nodes...
    else:
        mc.file(filePath, ir=True)
        mc.namespace(rm='shader', f=True, mnr=True)

        DG_modifier = OpenMaya.MDGModifier()

        iterator = OpenMaya.MItSelectionList(unuse_sg_nodes)
        mobject  = OpenMaya.MObject()
        while not iterator.isDone():
            iterator.getDependNode(mobject)

            for in_node in get_input_dg_nodes(mobject):
                DG_modifier.deleteNode(in_node)
            OpenMaya.MGlobal.deleteNode(mobject)

            iterator.next()

    return True





def create_scence_geo_data_dict(geometry_data, use_name_key=True):
    '''
    '''
    scene_geo_data = dict()
    for u, mobject in geometry_data:
        if use_name_key:
            n = re.search('\w+$', muuid.get_scene_object_name(mobject)).group()
            scene_geo_data.setdefault(n, list()).append(muuid.get_scene_object_name(mobject))
        else:
            scene_geo_data.setdefault(u, list()).append(muuid.get_scene_object_name(mobject))

    return scene_geo_data





def import_data_to_sel(filePath, by_name=True):
    '''
    '''
    selection = OpenMaya.MSelectionList()
    OpenMaya.MGlobal.getActiveSelectionList(selection)
    geometry_data = create_scence_geo_data_dict(muuid.get_scene_selection_uuid_data(selection), by_name)

    #-
    import_data_to_geo(filePath, geometry_data, by_name)






def import_data_to_all(filePath, by_name=True):
    '''
    '''
    geometry_data = create_scence_geo_data_dict(muuid.get_scene_geometry_uuid_data(True), by_name)

    #-
    import_data_to_geo(filePath, geometry_data, by_name)





def import_data_to_geo(filePath, geometry_data, by_name=True):
    '''
    '''
    if not os.path.isfile(filePath):
        return

    data_xml_doc = minidom.parse(filePath).documentElement
    mayaTool.startProgress(len(data_xml_doc.childNodes))
    for sg_xml_node in data_xml_doc.childNodes:
        set_SG_connet_geometry(sg_xml_node, geometry_data, by_name)
        mayaTool.moveProgress('Apply shading data from {0} - {1}'.format(sg_xml_node.getAttribute('name'), sg_xml_node.getAttribute('uuid')))
    mayaTool.endProgress()





def set_SG_connet_geometry(sg_xml_node, geometry_data, by_name=True):
    '''
    '''
    #- find shadingEngine node...
    sg_node_mobjects = list(muuid.get_objects_by_uuid(sg_xml_node.getAttribute('uuid'), OpenMaya.MFn.kShadingEngine, False))
    if not sg_node_mobjects:
        return
    sg_node_name     = muuid.get_scene_object_name(sg_node_mobjects[-1])

    #- find matched geometry...
    geometry_list    = list()
    for geo_xml_node in sg_xml_node.childNodes:
        if by_name:
            key = geo_xml_node.getAttribute('name').split('.')[0]
        else:
            key = geo_xml_node.getAttribute('uuid')

        geometry_list.extend([re.sub('\w+', m, geo_xml_node.getAttribute('name'), 1) for m in geometry_data.get(key, list())])

    if not geometry_list:
        return

    #- connect shadingEngine and geometry...
    mc.sets(geometry_list, e=True, forceElement=sg_node_name)





def merge_data_to_scene(shader_path_list, geoFilter=0, by_name=True):
    '''
    geoFilter:
        0 - all of geometry in scene
        1 - selection geometry
        2 - geometry in camera view
    '''
    selection = OpenMaya.MSelectionList()

    #- get all of geometry to selection list
    if geoFilter == 0:
        for mobject in itertools.chain(muuid.get_scene_dependency_nodes(OpenMaya.MFn.kNurbsSurface),
                                       muuid.get_scene_dependency_nodes(OpenMaya.MFn.kMesh),
                                       muuid.get_scene_dependency_nodes(OpenMaya.MFn.kSubdiv),
                                       muuid.get_scene_dag_parent_nodes(muuid.get_scene_dependency_nodes(OpenMaya.MFn.kNurbsCurve)),
                                       muuid.get_scene_dag_parent_nodes(muuid.get_scene_dependency_nodes(OpenMaya.MFn.kMesh)),
                                       muuid.get_scene_dag_parent_nodes(muuid.get_scene_dependency_nodes(OpenMaya.MFn.kSubdiv))):
            selection.add(mobject)

        for ais in pymel.core.ls(typ='aiStandIn'):
            ai_stand_mobj = ais.__apimobject__()

            selection.add(ai_stand_mobj)
            selection.add(OpenMaya.MFnDagNode(ai_stand_mobj).parent(0))


    #- get selection geometry to selection list
    elif geoFilter == 1:
        mc.select(hi=True)
        OpenMaya.MGlobal.getActiveSelectionList(selection)


    #- get camera view inside geometry to selection list
    elif geoFilter == 2:
        activeView = OpenMayaUI.M3dView.active3dView()
        OpenMaya.MGlobal.selectFromScreen(0, 0, activeView.portWidth(), activeView.portHeight(), OpenMaya.MGlobal.kReplaceList, OpenMaya.MGlobal.kSurfaceSelectMethod)

        mc.select(hi=True)
        OpenMaya.MGlobal.getActiveSelectionList(selection)


    if selection.isEmpty():
        return

    #- create geometry data dict
    geometry_data = create_scence_geo_data_dict(muuid.get_scene_selection_uuid_data(selection), by_name)
    scene_geo_uuid_data = create_scence_geo_data_dict(muuid.get_scene_selection_uuid_data(selection), False)

    #-
    for shaderPath in shader_path_list:
        if not os.path.isfile(shaderPath):
            continue

        dataPath = '{0}.xml'.format(os.path.splitext(shaderPath)[0])
        if not os.path.isfile(dataPath):
            continue

        #- read shader data file, check model is or not in scene...
        f_pointer = os.open(dataPath, os.O_RDONLY)
        f_mmap    = mmap.mmap(f_pointer, 0, access=mmap.ACCESS_READ)

        data_uuid_list = re.findall('[a-z0-9]{32}', f_mmap)
        data_uuid_list = scriptTool.arrayRemoveDuplicates(data_uuid_list)

        f_mmap.close()
        os.close(f_pointer)

        #-
        shader_data_can_use = False
        for uuid_code in data_uuid_list:
            if scene_geo_uuid_data.has_key(uuid_code):
                shader_data_can_use = True
                break
        if not shader_data_can_use:
            continue

        #- import shader file and shader data
        import_SG_nodes(shaderPath)
        import_data_to_geo(dataPath, geometry_data, by_name)
